from cryptography.fernet import Fernet
import base64
import hashlib


def get_fernet_key(password):
    # 将密码编码为字节
    password_bytes = password.encode()
    # 使用 SHA256 哈希算法生成 32 字节的密钥
    key = hashlib.sha256(password_bytes).digest()
    # 将密钥转换为 URL 安全的 Base64 编码
    fernet_key = base64.urlsafe_b64encode(key)
    return fernet_key


def encrypt_string(plaintext, password):
    fernet_key = get_fernet_key(password)
    f = Fernet(fernet_key)
    # 将明文编码为字节
    plaintext_bytes = plaintext.encode()
    # 加密字节
    encrypted = f.encrypt(plaintext_bytes)
    return encrypted


def decrypt_string(ciphertext, password):
    fernet_key = get_fernet_key(password)
    f = Fernet(fernet_key)
    # 解密字节
    decrypted = f.decrypt(ciphertext)
    # 将解密后的字节转换为字符串
    decrypted_text = decrypted.decode()
    return decrypted_text


# 示例用法
if __name__ == "__main__":
    # 中文密码
    password = "加密密码"
    # 要加密的字符串
    plaintext = "需要加密的内容"

    # 加密
    encrypted = encrypt_string(plaintext, password)
    print(f"加密后的内容: {encrypted}")

    # 解密
    decrypted = decrypt_string(encrypted, password)
    print(f"解密后的内容: {decrypted}")
